import random as rand
import numpy as np


def weight_from_edge(edge):
    return edge[2]


def random_color():
    rand_color = (
        int(rand.random() * 255),
        int(rand.random() * 255),
        int(rand.random() * 255),
    )
    return rand_color


def generate_color_scheme(arr):
    """
    Generate color scheme for pseudo colored label map

    Parameters
    ----------
    arr : ndarray
        label map numpy array

    Returns
    -------
    color_scheme : list
        random color scheme list
    """
    height, width = arr.shape
    color_scheme = [random_color() for i in range(width * height)]

    return color_scheme


def generate_label_map(ufset, width, height):
    """Generate label map from the union find set, which is the segmentation graph from the FH superpixel algorithm

    Parameters
    ----------
    ufset : instance of OptimizedUnionFind
        the union find set, which is the segmentation graph from the FH superpixel algorithm
    width : int
        image width
    height : int
        image height

    Returns
    -------
    save_label_map : ndarray
        a superpixel segmentation label map
    """
    save_label_map = np.zeros((height, width), np.uint32)

    color_set = set()
    for y in range(height):
        for x in range(width):
            color_idx = ufset.find(y * width + x)
            save_label_map[y, x] = color_idx
            color_set.add(color_idx)

    return save_label_map


def generate_label_image(ufset, width, height):
    """Generate pseudo color segmentation label image from the union find set,
    which is the segmentation graph from the FH superpixel algorithm

    Parameters
    ----------
    ufset : instance of OptimizedUnionFind
        the union find set, which is the segmentation graph from the FH superpixel algorithm
    width : int
        image width
    height : int
        image height

    Returns
    -------
    save_img : ndarray
        a superpixel segmentation label map
    """

    color = [random_color() for i in range(width * height)]

    save_img = np.zeros((height, width, 3), np.uint8)

    for y in range(height):
        for x in range(width):
            color_idx = ufset.find(y * width + x)
            save_img[y, x] = color[color_idx]

    return save_img


def generate_label_image_for_array(arr, color_scheme=None):
    """Generate pseudo color segmentation label image from a ndarray label map

    Parameters
    ----------
    arr : ndarray
        a ndarray label map [h x w]
    color_scheme : list, optional
        random color scheme, by default None

    Returns
    -------
    save_img : ndarray
        the ndarray pseudo color map [h x w x 3] ready for saving
    """

    height, width = arr.shape

    if color_scheme is None:
        color_scheme = [random_color() for i in range(width * height)]

    save_img = np.zeros((height, width, 3), np.uint8)

    for y in range(height):
        for x in range(width):
            color_idx = arr[y, x]
            save_img[y, x] = color_scheme[color_idx]

    return save_img
